package agent.configuration.rag;

import agent.configuration.rag.util.SplitterUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.model.embedding.DisabledEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.math.BigDecimal;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocument;

@RequiredArgsConstructor
@Configuration
public class EmbeddingConfig {


    private final SplitterUtil splitterUtil;

    @Value("${Embedding.documentPath}")
    private String documentPath;

    @Value("${Embedding.apiKey}")
    private String embeddingApiKey;


    @Value("${Embedding.modelName}")
    private String embeddingModelName;

    @Bean
    public Document Document(){

        return  loadDocument(documentPath, new TextDocumentParser());

    }

    @Bean
    public EmbeddingModel DashScore(){
        return new DisabledEmbeddingModel();
    }

    @Bean(name = "qwenEmbeddingModel")
    public EmbeddingModel BgeSmallEnV15QuantizedEmbeddingModel(){
        return  new QwenEmbeddingModel(embeddingApiKey,embeddingModelName);
    }

    @Bean(name = "inMemoryEmbeddingStore")
    public EmbeddingStore<TextSegment> InMemoryEmbeddingStore(EmbeddingModel qwenEmbeddingModel, Document document){


        InMemoryEmbeddingStore<TextSegment> textSegmentEmbeddingStore = new InMemoryEmbeddingStore<>();

        makingIngest(qwenEmbeddingModel, textSegmentEmbeddingStore,document);

        return textSegmentEmbeddingStore;
    }

    private  void makingIngest(EmbeddingModel embeddingModel, InMemoryEmbeddingStore<TextSegment> textSegmentEmbeddingStore,Document document) {
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(DocumentSplitters.recursive(300, 0))

                .embeddingModel(embeddingModel)
                .embeddingStore(textSegmentEmbeddingStore)
                .build();
        //    文档已被分割为片段、生成向量，并存储到 embeddingStore 中
        ingestor.ingest(document);
    }

    @Bean
    public ContentRetriever getEmbeddingContentRetriever(EmbeddingModel qwenEmbeddingModel,EmbeddingStore<TextSegment> inMemoryEmbeddingStore){

        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(inMemoryEmbeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(2) //maxResults（最大结果数）
                .minScore(0.6) //  minScore（最小相似度分数）
                .build();

    }

}
